#encoding=utf8
import numpy as np
import pickle
import os
import sys
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.naive_bayes import GaussianNB


def load_dataset(file_name):
    '''
    从文件读入数据集
    被多处调用，请勿删除或改动本函数！！！
    '''
    try:
        with open(file_name, 'rb') as f:
            raw_dataset = pickle.load(f)
    except FileNotFoundError:
        print(f"错误: 文件 {file_name} 未找到。请确保文件路径正确。")
        return None, None
    
    try:
        example_image = raw_dataset[0][0]
    except KeyError:
        print("错误: 数据集格式不正确，无法找到类别0的数据。")
        return None, None
    except TypeError:
        print("错误: 数据集格式不正确，类别0的数据不是列表或数组。")
        return None, None

    dataset = np.empty((0, example_image.size))
    labels = np.empty((0, 1))
    
    total_images = 0
    for i_class in raw_dataset.keys():
        images_list = raw_dataset.get(i_class, [])
        if not isinstance(images_list, list) or len(images_list) == 0:
            continue
        for image in images_list:
            features = image.flatten() / 255.0  # 像素归一化到 [0,1]
            dataset = np.vstack((dataset, features))
            labels = np.vstack((labels, i_class))
            total_images += 1
            
    print(f"成功加载 {total_images} 张手写数字图片。")
    return dataset, labels


class NBClassifier:
    def __init__(self, train_dataset=None, train_labels=None):
        self.model = None
        if train_dataset is None or train_labels is None:
            self.train_dataset, self.train_labels = load_dataset('./step1/input/training_dataset.pkl')
            if self.train_dataset is None:
                raise RuntimeError("训练数据集加载失败，无法继续。")
        else:
            self.train_dataset = train_dataset
            self.train_labels = train_labels

    def train(self):
        print("=== 正在训练 GaussianNB 模型... ===")
        self.model = GaussianNB()
        train_labels = self.train_labels.ravel()
        self.model.fit(self.train_dataset, train_labels)
        print("=== 模型训练完成 ===")

    def predict(self, test_dataset):
        '''
        输入：测试数据 test_dataset: 形状为(500, 784)的ndarray
        输出：预测结果 predicted_labels: 形状为(500, )的ndarray
        '''
        predicted_labels = self.model.predict(test_dataset)
        return predicted_labels


def calculate_accuracy(file_name, classifier):
    test_dataset, test_labels = load_dataset(file_name)
    if test_dataset is None:
        return 0
    random_indices = np.random.permutation(test_dataset.shape[0])
    test_dataset = test_dataset[random_indices,:]
    test_labels = test_labels[random_indices,:]
    predicted_labels = classifier.predict(test_dataset)
    if isinstance(predicted_labels, np.ndarray):
        if predicted_labels.size != test_labels.size:
            print('错误：输出的标签数量与测试集大小不一致')
            accuracy = 0
        else:
            accuracy = np.mean(predicted_labels.flatten()==test_labels.flatten())
    else:
        print('错误：输出格式有误，必须为ndarray格式')
        accuracy = 0
    return accuracy


if __name__ == '__main__':
    # ========== 少量训练数据实验 ==========
    os.makedirs("./step1/output/nb2_less", exist_ok=True)
    result_path = "./step1/output/nb2_less/result.md"
    figure_path = "./step1/output/nb2_less/result.png"

    full_train_dataset, full_train_labels = load_dataset('./step1/input/training_dataset.pkl')
    results = []
    test_file = './step1/input/test_dataset_clean.pkl'

    for percent in range(100, 0, -10):
        sample_size = int(full_train_dataset.shape[0] * percent / 100)
        indices = np.random.choice(full_train_dataset.shape[0], sample_size, replace=False)
        reduced_train_dataset = full_train_dataset[indices]
        reduced_train_labels = full_train_labels[indices]

        classifier = NBClassifier(train_dataset=reduced_train_dataset, train_labels=reduced_train_labels)
        classifier.train()
        acc = calculate_accuracy(test_file, classifier)
        results.append((percent, acc))
        print(f"使用 {percent}% 训练数据，测试集准确率: {acc:.4f}")

    # ========== 保存结果到 markdown ==========
    with open(result_path, "w", encoding="utf-8") as f:
        f.write("# GaussianNB 少量训练数据实验结果\n\n")
        f.write("| 使用训练数据比例 | 测试集准确率 |\n")
        f.write("|----------------|--------------|\n")
        for percent, acc in results:
            f.write(f"| {percent}% | {acc:.4f} |\n")

    # ========== 画图 ==========
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    percents = [p for p, _ in results]
    accuracies = [a for _, a in results]
    plt.figure(figsize=(8, 6))
    plt.plot(percents, accuracies, marker='o', linestyle='-', color='g', label="测试集准确率")
    plt.title("GaussianNB 不同训练数据比例的效果")
    plt.xlabel("训练数据比例 (%)")
    plt.ylabel("测试集准确率")
    plt.xticks(percents)
    plt.ylim(0.0, 1.0)  # 保持完整区间
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.legend()
    plt.savefig(figure_path)
    plt.close()

    print(f"\n结果已保存到 {result_path}")
    print(f"折线图已保存到 {figure_path}")
